#!/usr/bin/python
#
#

import sys
import re
from Bio.Blast import NCBIXML

xml_file = sys.argv[1]
xf = open(xml_file, "rU")
r = NCBIXML.parse(xf)
results_rec = r.next()

cc_file = sys.argv[2]
ccf = open(cc_file, "rU")
ccl = ccf.readlines()

m1 = re.compile('.* (COG\d+), (\w+),.*')
m2 = re.compile('(COG\d+), (\w+|\w+-\w+|\w+-\w+-\w+|\{.*\}),.*') #2nd group matches regular words or hyphenated words
m3 = re.compile('(.*).cog.rpsblast.xml')
m4 = re.compile('\[(.*)\]')

desc_obj_list = []

cds_id = m3.match(xml_file).group(1)

max_evalue = 0.01 #0.01 matches well with SMART EMBL's domain search results

cog_ids = []
cog_descs = []
query_starts = []
query_stops = []
evalues = []

#cogslist = []
cogs_dict = {}

for line in ccl:
    l = line.split('\t')
    cog_cat = m4.match(l[0]).group(1)
    cog_id = l[1]
    cog_desc = l[2].rstrip()
    #cog_info = (cog_id, cog_cat, cog_desc)
    #cogslist.append(cog_info)
    cogs_dict[cog_id] = (cog_cat, cog_desc)

if len(results_rec.alignments) == 0:
    #print cds_id + "\t" + "No COG Hit"
    pass
else:
    for alignment in results_rec.alignments:
        p = m2.match(alignment.hit_def)
        pfam_id = p.group(1)
        pfam_desc = p.group(2)
        if len(alignment.hsps) == 1:
            for hsp in alignment.hsps:
                if hsp.expect < max_evalue:
                    #print cds_id + "\t" + pfam_id + "\t" + pfam_desc + "\t" + str(hsp.query_start) + "\t" + str(hsp.query_end) + "\t" +str(hsp.expect)
                    cog_ids.append(pfam_id)
                    cog_descs.append(pfam_desc)
                    query_starts.append(hsp.query_start)
                    query_stops.append(hsp.query_end)
                    evalues.append(hsp.expect)
        else:
            hsps_list = []
            for hsp in alignment.hsps:
                hsps_list.append(hsp)
            for h in hsps_list:
                if h.expect < max_evalue:
                    cog_ids.append(pfam_id)
                    cog_descs.append(pfam_desc)
                    query_starts.append(h.query_start)
                    query_stops.append(h.query_end)
                    evalues.append(h.expect)

if len(evalues) == 0:
    print cds_id + "\t" + "-" + "\t" + "-" + "\t" + "-"
else:
    smallest = min(evalues)
    min_eindex = evalues.index(smallest)
    print cds_id + "\t" + cog_ids[min_eindex] + "\t" + cogs_dict[cog_ids[min_eindex]][0] + "\t" + str(evalues[min_eindex])
    # Note: cogs_dict[cog_ids[min_eindex]][0] -> extracts the first tuple from dictionary key, value pairs.
xf.close()